import os
import cv2
import numpy as np
from flask import Flask, render_template, request, jsonify, Response
from PIL import Image, ImageDraw, ImageFont
from io import BytesIO
import requests
from ultralytics import YOLO
import base64
import json
from datetime import datetime
import threading
import queue
import time

app = Flask(__name__)

# 配置
UPLOAD_FOLDER = 'static/uploads'
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif'}
MODEL_PATH = 'models/best.pt'  # YOLOv8模型路径
FONT_PATH = 'static/fonts/SimHei.ttf'  # 中文字体路径，需要确保存在

# 确保上传文件夹和字体目录存在
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(os.path.dirname(FONT_PATH), exist_ok=True)

# 如果字体不存在，尝试使用系统字体
if not os.path.exists(FONT_PATH):
    # 尝试在Windows系统上找到常见的中文字体
    windows_font_path = "C:/Windows/Fonts/simhei.ttf"
    if os.path.exists(windows_font_path):
        FONT_PATH = windows_font_path
    else:
        print("警告：未找到中文字体文件，标签可能显示为乱码")

app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 限制上传文件大小为16MB

# 水果信息数据库
FRUITS_INFO = {
    'apple': {
        'name': '苹果',
        'description': '苹果是一种常见的水果，富含膳食纤维和维生素。',
        'nutrition': '每100克含有：\n• 热量：52千卡\n• 碳水化合物：14克\n• 膳食纤维：2.4克\n• 维生素C：4.6毫克\n• 钾：107毫克',
        'benefits': '功效作用：\n• 促进肠道健康\n• 调节血糖\n• 增强免疫力\n• 改善心血管健康\n• 帮助减重',
        'precautions': '食用禁忌：\n• 空腹不宜食用\n• 胃寒人群慎食\n• 术后病人不宜多食\n• 牙齿敏感者注意'
    },
    'banana': {
        'name': '香蕉',
        'description': '香蕉是一种富含钾和维生素B6的水果，能够提供持久能量。',
        'nutrition': '每100克含有：\n• 热量：89千卡\n• 碳水化合物：23克\n• 膳食纤维：2.6克\n• 钾：358毫克\n• 维生素B6：0.4毫克',
        'benefits': '功效作用：\n• 补充能量\n• 改善心情\n• 缓解肌肉疲劳\n• 促进消化\n• 帮助睡眠',
        'precautions': '食用禁忌：\n• 糖尿病人慎食\n• 肾功能不全者少食\n• 过量可能导致便秘\n• 空腹不宜食用'
    },
    'orange': {
        'name': '橙子',
        'description': '橙子富含维生素C和类黄酮，是增强免疫力的理想水果。',
        'nutrition': '每100克含有：\n• 热量：47千卡\n• 碳水化合物：12克\n• 维生素C：53.2毫克\n• 叶酸：30微克\n• 钾：181毫克',
        'benefits': '功效作用：\n• 增强免疫力\n• 美白护肤\n• 预防感冒\n• 促进胶原蛋白合成\n• 改善血液循环',
        'precautions': '食用禁忌：\n• 胃炎患者慎食\n• 口腔溃疡者少食\n• 孕妇不宜过量\n• 空腹不宜食用'
    },
    'strawberry': {
        'name': '草莓',
        'description': '草莓富含抗氧化物质和维生素C，有助于保护心脏健康。',
        'nutrition': '每100克含有：\n• 热量：32千卡\n• 碳水化合物：7.7克\n• 维生素C：58.8毫克\n• 叶酸：24微克\n• 钾：153毫克',
        'benefits': '功效作用：\n• 抗氧化\n• 美白护肤\n• 保护视力\n• 预防心血管疾病\n• 改善记忆力',
        'precautions': '食用禁忌：\n• 过敏体质者慎食\n• 肾结石患者少食\n• 需充分清洗\n• 糖尿病人控制量'
    },
    'grape': {
        'name': '葡萄',
        'description': '葡萄含有丰富的白藜芦醇和多酚类物质，具有很强的抗氧化作用。',
        'nutrition': '每100克含有：\n• 热量：69千卡\n• 碳水化合物：18克\n• 维生素C：3.2毫克\n• 钾：191毫克\n• 铁：0.36毫克',
        'benefits': '功效作用：\n• 抗衰老\n• 保护心血管\n• 改善睡眠\n• 提高免疫力\n• 预防癌症',
        'precautions': '食用禁忌：\n• 糖尿病人慎食\n• 牙齿敏感者少食\n• 空腹不宜食用\n• 发烧期间避免食用'
    },
    'kiwi': {
        'name': '猕猴桃',
        'description': '猕猴桃富含维生素C、维生素E和膳食纤维，是营养价值极高的水果。',
        'nutrition': '每100克含有：\n• 热量：61千卡\n• 碳水化合物：15克\n• 维生素C：92.7毫克\n• 维生素E：1.5毫克\n• 钾：312毫克',
        'benefits': '功效作用：\n• 促进消化\n• 改善睡眠\n• 增强免疫力\n• 美容养颜\n• 预防便秘',
        'precautions': '食用禁忌：\n• 过敏体质者慎食\n• 肾结石患者少食\n• 空腹不宜食用\n• 与乳制品间隔食用'
    },
    'lemon': {
        'name': '柠檬',
        'description': '柠檬富含维生素C和柠檬酸，具有很好的美白和解毒作用。',
        'nutrition': '每100克含有：\n• 热量：29千卡\n• 碳水化合物：9.3克\n• 维生素C：53毫克\n• 钾：138毫克\n• 柠檬酸：5克',
        'benefits': '功效作用：\n• 美白护肤\n• 解毒排毒\n• 提高免疫力\n• 促进新陈代谢\n• 帮助减重',
        'precautions': '食用禁忌：\n• 胃病患者慎食\n• 牙齿敏感者少食\n• 空腹不宜食用\n• 过量可能损害胃黏膜'
    },
    'mango': {
        'name': '芒果',
        'description': '芒果富含维生素A、维生素C和膳食纤维，是一种营养丰富的热带水果。',
        'nutrition': '每100克含有：\n• 热量：60千卡\n• 碳水化合物：15克\n• 维生素A：54微克\n• 维生素C：36.4毫克\n• 膳食纤维：1.6克',
        'benefits': '功效作用：\n• 改善视力\n• 增强免疫力\n• 促进消化\n• 美容养颜\n• 预防贫血',
        'precautions': '食用禁忌：\n• 糖尿病人慎食\n• 过敏体质者少食\n• 发烧期间避免\n• 空腹不宜食用'
    },
    'pineapple': {
        'name': '菠萝',
        'description': '菠萝含有丰富的维生素C和菠萝蛋白酶，具有很好的消化功能。',
        'nutrition': '每100克含有：\n• 热量：50千卡\n• 碳水化合物：13克\n• 维生素C：47.8毫克\n• 锰：0.927毫克\n• 膳食纤维：1.4克',
        'benefits': '功效作用：\n• 促进消化\n• 消除疲劳\n• 美容养颜\n• 预防感冒\n• 消除水肿',
        'precautions': '食用禁忌：\n• 胃病患者慎食\n• 孕妇慎食\n• 空腹不宜食用\n• 与牛奶间隔食用'
    },
    'watermelon': {
        'name': '西瓜',
        'description': '西瓜水分含量高，富含番茄红素和氨基酸，具有很好的解暑功效。',
        'nutrition': '每100克含有：\n• 热量：30千卡\n• 碳水化合物：8克\n• 维生素C：8.1毫克\n• 钾：112毫克\n• 番茄红素：4.1毫克',
        'benefits': '功效作用：\n• 解暑降温\n• 利尿消肿\n• 预防中暑\n• 美容养颜\n• 降血压',
        'precautions': '食用禁忌：\n• 糖尿病人慎食\n• 脾胃虚寒者少食\n• 空腹不宜食用\n• 不宜过量食用'
    }
}

# 加载YOLOv8模型
def load_model():
    try:
        model = YOLO(MODEL_PATH)
        print("模型加载成功！")
        return model
    except Exception as e:
        print(f"模型加载失败: {e}")
        return None

# 检查文件扩展名是否允许
def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

# 处理图像并进行预测
def process_image(image_path, model):
    try:
        # 使用YOLOv8进行预测
        results = model(image_path)
        
        # 获取预测结果
        result = results[0]
        
        # 提取检测到的对象、置信度和边界框
        detections = []
        
        for box in result.boxes:
            cls_id = int(box.cls.item())
            cls_name = result.names[cls_id]
            confidence = float(box.conf.item())
            
            # 只返回置信度大于0.5的结果
            if confidence > 0.5:
                x1, y1, x2, y2 = box.xyxy[0].tolist()
                detections.append({
                    'class': cls_name,
                    'confidence': confidence,
                    'bbox': [x1, y1, x2, y2]
                })
        
        # 读取原始图像并转为PIL图像对象
        cv_img = cv2.imread(image_path)
        pil_img = Image.fromarray(cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB))
        draw = ImageDraw.Draw(pil_img)
        
        # 加载中文字体
        try:
            font = ImageFont.truetype(FONT_PATH, 20)  # 字体大小适当调小
        except IOError:
            # 如果找不到指定的字体，使用默认字体
            font = ImageFont.load_default()
            print("警告：无法加载中文字体，使用默认字体")
        
        for det in detections:
            x1, y1, x2, y2 = map(int, det['bbox'])
            cls_name = det['class']
            confidence = det['confidence']
            
            # 获取中文名称
            fruit_info = get_fruit_info(cls_name)
            chinese_name = fruit_info['name']
            
            # 绘制边界框
            draw.rectangle((x1, y1, x2, y2), outline=(0, 255, 0), width=2)
            
            # 准备标签文本
            label = f"{chinese_name}: {confidence:.2f}"
            
            # 计算文本大小
            text_bbox = draw.textbbox((0, 0), label, font=font)
            text_width = text_bbox[2] - text_bbox[0]
            text_height = text_bbox[3] - text_bbox[1]
            
            # 计算标签位置 - 确保标签在图像内
            box_width = x2 - x1
            
            # 如果文本宽度大于边界框宽度，则将标签放在边界框上方
            if text_width > box_width:
                label_x = max(0, x1 - (text_width - box_width) // 2)
                label_y = max(0, y1 - text_height - 5)
            else:
                # 否则将标签放在边界框内顶部
                label_x = x1 + 5
                label_y = y1 + 5
            
            # 确保标签不会超出图像边界
            img_width, img_height = pil_img.size
            if label_x + text_width > img_width:
                label_x = max(0, img_width - text_width - 5)
            
            # 绘制文本背景框
            draw.rectangle(
                (label_x, label_y, label_x + text_width, label_y + text_height),
                fill=(0, 255, 0)
            )
            
            # 添加标签文本
            draw.text((label_x, label_y), label, fill=(0, 0, 0), font=font)
        
        # 将PIL图像转回OpenCV格式并保存
        cv_output = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
        # 修改路径生成方式，只替换最后一个点
        base, ext = os.path.splitext(image_path)
        output_path = base + '_detected' + ext
        cv2.imwrite(output_path, cv_output)
        
        # 获取水果信息
        results = []
        for det in detections:
            fruit_class = det['class']
            confidence = det['confidence']
            fruit_info = get_fruit_info(fruit_class)
            
            results.append({
                'class': fruit_class,
                'chinese_name': fruit_info['name'],
                'confidence': confidence,
                'description': fruit_info['description'],
                'nutrition': fruit_info['nutrition'],
                'benefits': fruit_info['benefits'],
                'precautions': fruit_info['precautions']
            })
        
        # 确保返回正确的Web可访问路径
        web_path = get_web_path(output_path)
        print(f"图像处理完成，输出路径: {output_path}, Web路径: {web_path}")
        
        return {
            'success': True,
            'results': results,
            'image_path': web_path
        }
    except Exception as e:
        print(f"处理图像时出错: {e}")
        return {
            'success': False,
            'error': str(e)
        }

# 获取水果信息
def get_fruit_info(fruit_class):
    # 将YOLOv8的类名转换为我们的水果信息数据库键
    fruit_class = fruit_class.lower()
    return FRUITS_INFO.get(fruit_class, {
        'name': fruit_class,
        'description': '暂无描述',
        'nutrition': '暂无营养信息'
    })

# 主页
@app.route('/')
def index():
    return render_template('index.html')

# 处理上传的图片
@app.route('/detect', methods=['POST'])
def detect():
    if 'file' not in request.files:
        return jsonify({'error': '没有上传文件'}), 400
    
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': '没有选择文件'}), 400
    
    if file:
        # 生成唯一文件名
        filename = f"{datetime.now().strftime('%Y%m%d_%H%M%S')}_{file.filename}"
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(filepath)
        
        # 加载模型
        model = load_model()
        if model is None:
            return jsonify({'error': '模型加载失败'})
        
        # 处理图像并获取预测结果
        result = process_image(filepath, model)
        
        # 如果没有检测到水果
        if 'error' in result:
            return jsonify(result)
        
        # 返回结果
        return jsonify(result)
    
    return jsonify({'error': '不支持的文件类型'})

# 处理URL图片
@app.route('/url', methods=['POST'])
def process_url():
    data = request.get_json()
    
    if not data or 'url' not in data:
        return jsonify({'error': '未提供URL'})
    
    image_url = data['url']
    
    try:
        # 下载图片
        response = requests.get(image_url)
        img = Image.open(BytesIO(response.content))
        
        # 保存图片
        filename = os.path.join(app.config['UPLOAD_FOLDER'], 'url_image.jpg')
        img.save(filename)
        
        # 加载模型
        model = load_model()
        if model is None:
            return jsonify({'error': '模型加载失败'})
        
        # 处理图像并获取预测结果
        result = process_image(filename, model)
        
        # 如果没有检测到水果
        if 'error' in result:
            return jsonify(result)
        
        # 返回结果
        return jsonify(result)
    
    except Exception as e:
        return jsonify({'error': f'处理URL图片时出错: {str(e)}'})

# 全局变量用于存储摄像头状态
camera = None
camera_active = False
frame_queue = queue.Queue(maxsize=10)

# 辅助函数：使用PIL绘制带中文标签的帧
def draw_chinese_labels(frame, detections):
    # 转换为PIL图像
    pil_img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(pil_img)
    img_width, img_height = pil_img.size
    
    # 加载中文字体
    try:
        font = ImageFont.truetype(FONT_PATH, 20)
    except IOError:
        font = ImageFont.load_default()
        print("警告：无法加载中文字体，使用默认字体")
    
    for det in detections:
        x1, y1, x2, y2 = map(int, det['bbox'])
        cls_name = det['class']
        confidence = det['confidence']
        
        # 获取中文名称
        fruit_info = get_fruit_info(cls_name)
        chinese_name = fruit_info['name']
        
        # 绘制边界框
        draw.rectangle((x1, y1, x2, y2), outline=(0, 255, 0), width=2)
        
        # 准备标签文本
        label = f"{chinese_name}: {confidence:.2f}"
        
        # 计算文本大小
        text_bbox = draw.textbbox((0, 0), label, font=font)
        text_width = text_bbox[2] - text_bbox[0]
        text_height = text_bbox[3] - text_bbox[1]
        
        # 计算标签位置 - 确保标签在图像内
        box_width = x2 - x1
        
        # 如果文本宽度大于边界框宽度，则将标签放在边界框上方
        if text_width > box_width:
            label_x = max(0, x1 - (text_width - box_width) // 2)
            label_y = max(0, y1 - text_height - 5)
        else:
            # 否则将标签放在边界框内顶部
            label_x = x1 + 5
            label_y = y1 + 5
        
        # 确保标签不会超出图像边界
        if label_x + text_width > img_width:
            label_x = max(0, img_width - text_width - 5)
        
        # 绘制文本背景框
        draw.rectangle(
            (label_x, label_y, label_x + text_width, label_y + text_height),
            fill=(0, 255, 0)
        )
        
        # 添加标签文本
        draw.text((label_x, label_y), label, fill=(0, 0, 0), font=font)
    
    # 转回OpenCV格式
    return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)

def generate_frames():
    global camera, camera_active
    model = load_model()
    
    while camera_active:
        success, frame = camera.read()
        if not success:
            print("无法从摄像头读取帧")
            time.sleep(0.1)
            continue
        
        try:
            # 运行模型推理
            results = model(frame)
            result = results[0]
            
            # 提取检测结果
            detections = []
            
            for box in result.boxes:
                cls_id = int(box.cls.item())
                cls_name = result.names[cls_id]
                confidence = float(box.conf.item())
                
                if confidence > 0.5:
                    x1, y1, x2, y2 = box.xyxy[0].tolist()
                    detections.append({
                        'class': cls_name,
                        'confidence': confidence,
                        'bbox': [x1, y1, x2, y2]
                    })
            
            # 使用中文标签绘制
            if detections:
                frame = draw_chinese_labels(frame, detections)
            
            # 将帧转换为JPEG格式
            ret, buffer = cv2.imencode('.jpg', frame)
            frame = buffer.tobytes()
            
            # 将检测结果放入队列
            try:
                detection_results = [
                    {
                        'class': det['class'],
                        'confidence': det['confidence']
                    }
                    for det in detections
                ]
                
                frame_queue.put_nowait({
                    'frame': frame,
                    'detections': detection_results
                })
                
                # 打印检测到的物体信息
                if detection_results:
                    print(f"检测到的物体: {[d['class'] for d in detection_results]}")
                
            except queue.Full:
                frame_queue.get()  # 如果队列满了，移除最旧的帧
                frame_queue.put({
                    'frame': frame,
                    'detections': detection_results
                })
            
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')
        
        except Exception as e:
            print(f"处理摄像头帧时出错: {e}")
            import traceback
            traceback.print_exc()
            time.sleep(0.1)

@app.route('/video_feed')
def video_feed():
    return Response(generate_frames(),
                   mimetype='multipart/x-mixed-replace; boundary=frame')

@app.route('/start_camera', methods=['POST'])
def start_camera():
    global camera, camera_active
    try:
        if camera is None:
            camera = cv2.VideoCapture(0)  # 使用默认摄像头
            if not camera.isOpened():
                return jsonify({'error': '无法打开摄像头'}), 500
        
        # 设置摄像头参数
        camera.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        camera.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        
        camera_active = True
        print("摄像头已启动")
        return jsonify({'success': True})
    except Exception as e:
        print(f"启动摄像头时出错: {e}")
        import traceback
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500

@app.route('/stop_camera', methods=['POST'])
def stop_camera():
    global camera, camera_active
    try:
        camera_active = False
        if camera is not None:
            camera.release()
            camera = None
        print("摄像头已停止")
        return jsonify({'success': True})
    except Exception as e:
        print(f"停止摄像头时出错: {e}")
        import traceback
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500

@app.route('/get_latest_detections')
def get_latest_detections():
    try:
        if not frame_queue.empty():
            data = frame_queue.get()
            detections = data['detections']
            results = []
            
            for det in detections:
                fruit_class = det['class']
                confidence = det['confidence']
                fruit_info = get_fruit_info(fruit_class)
                
                results.append({
                    'class': fruit_class,
                    'chinese_name': fruit_info['name'],
                    'confidence': confidence,
                    'description': fruit_info['description'],
                    'nutrition': fruit_info['nutrition'],
                    'benefits': fruit_info['benefits'],
                    'precautions': fruit_info['precautions']
                })
            
            print(f"返回检测结果: {[r['chinese_name'] for r in results]}")
            return jsonify({'results': results})
        
        return jsonify({'results': []})
    except Exception as e:
        print(f"获取最新检测时出错: {e}")
        import traceback
        traceback.print_exc()
        return jsonify({'error': str(e)}), 500

# 处理视频路径，确保可以正确访问
def get_web_path(file_path):
    # 移除绝对路径前缀，只保留相对于项目的路径
    rel_path = file_path
    # 如果路径包含完整的磁盘路径，只保留static之后的部分
    if 'static' in file_path:
        rel_path = file_path[file_path.find('static'):]
    
    # 将反斜杠替换为正斜杠
    rel_path = rel_path.replace('\\', '/')
    
    # 确保路径不以/开头，因为这可能导致浏览器将其解析为绝对路径
    if rel_path.startswith('/'):
        rel_path = rel_path[1:]
    
    # 打印路径信息，便于调试    
    print(f"原始路径: {file_path}, 处理后路径: {rel_path}")
    return rel_path

def process_video(video_path, model):
    try:
        # 打开视频文件
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            print(f"错误：无法打开视频文件: {video_path}")
            return {'success': False, 'error': '无法打开视频文件'}
        
        # 获取视频信息
        fps = int(cap.get(cv2.CAP_PROP_FPS))
        frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        
        print(f"视频信息 - FPS: {fps}, 分辨率: {frame_width}x{frame_height}, 总帧数: {total_frames}")
        
        if fps <= 0 or frame_width <= 0 or frame_height <= 0:
            print(f"错误：无效的视频属性 - FPS: {fps}, 分辨率: {frame_width}x{frame_height}")
            return {'success': False, 'error': '无效的视频属性'}
        
        # 创建输出视频文件 - 使用更好的设置生成MP4视频
        output_filename = os.path.splitext(video_path)[0] + '_detected.mp4'
        
        # 尝试使用更兼容的编码器
        if os.name == 'nt':  # Windows系统
            # 在Windows上尝试使用MJPG或WMV2格式，通常这两种比较兼容
            try:
                fourcc = cv2.VideoWriter_fourcc(*'MJPG')
                print("使用MJPG编码器")
            except:
                fourcc = cv2.VideoWriter_fourcc(*'WMV2')
                print("使用WMV2编码器")
        else:
            # 在其他系统尝试使用MP4V格式
            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
            print("使用mp4v编码器")
        
        # 降低输出视频的尺寸以减小文件大小，提高兼容性
        output_width = min(frame_width, 800)
        output_height = int(frame_height * (output_width / frame_width))
        
        out = cv2.VideoWriter(output_filename, fourcc, fps, (output_width, output_height))
        
        if not out.isOpened():
            print(f"错误：无法创建输出视频文件: {output_filename}")
            return {'success': False, 'error': '无法创建输出视频文件'}
        
        # 为GIF预览准备的帧
        gif_frames = []
        gif_frame_interval = max(1, int(total_frames / 10))  # 最多采样10帧
        gif_output = os.path.splitext(video_path)[0] + '_preview.gif'
        
        all_detections = []
        frame_count = 0
        last_progress = -1
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            # 进度处理
            if total_frames > 0:
                progress = int((frame_count / total_frames) * 100)
                if progress != last_progress and progress % 10 == 0:  # 每10%显示一次
                    print(f"视频处理进度: {progress}%")
                    last_progress = progress
            
            # 调整输出帧的大小
            frame = cv2.resize(frame, (output_width, output_height))
            
            # 每隔几帧进行一次检测，根据视频长度调整检测间隔
            skip_frames = 3
            if total_frames > 1000:  # 对于较长的视频，增加跳帧数以加快处理速度
                skip_frames = 5
            
            if frame_count % skip_frames == 0:
                try:
                    results = model(frame)
                    result = results[0]
                    
                    frame_detections = []
                    detections = []
                    
                    for box in result.boxes:
                        cls_id = int(box.cls.item())
                        cls_name = result.names[cls_id]
                        confidence = float(box.conf.item())
                        
                        if confidence > 0.5:
                            x1, y1, x2, y2 = box.xyxy[0].tolist()
                            # 缩放检测框以匹配调整后的帧尺寸
                            x1 = int(x1 * (output_width / frame_width))
                            y1 = int(y1 * (output_height / frame_height))
                            x2 = int(x2 * (output_width / frame_width))
                            y2 = int(y2 * (output_height / frame_height))
                            
                            detections.append({
                                'class': cls_name,
                                'confidence': confidence,
                                'bbox': [x1, y1, x2, y2]
                            })
                            
                            frame_detections.append({
                                'class': cls_name,
                                'confidence': confidence
                            })
                    
                    # 使用中文标签绘制
                    if detections:
                        frame = draw_chinese_labels(frame, detections)
                    
                    all_detections.extend(frame_detections)
                except Exception as e:
                    print(f"处理第 {frame_count} 帧时出错: {e}")
                    # 继续处理，不中断整个过程
            
            # 写入输出视频
            out.write(frame)
            
            # 收集GIF预览的帧
            if frame_count % gif_frame_interval == 0:
                # 转换为PIL图像并添加到GIF帧列表
                gif_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                gif_frames.append(Image.fromarray(gif_frame))
            
            frame_count += 1
            
            # 安全检查，防止无限循环
            if frame_count > total_frames * 1.5 and total_frames > 0:
                print(f"警告：已处理 {frame_count} 帧，超过预期帧数 {total_frames}，停止处理")
                break
        
        # 释放资源
        cap.release()
        out.release()
        
        if frame_count == 0:
            print("错误：未能处理任何视频帧")
            return {'success': False, 'error': '视频处理失败，未能处理任何帧'}
        
        print(f"视频处理完成，共处理 {frame_count} 帧")
        
        # 生成GIF预览
        gif_created = False
        if gif_frames:
            try:
                print(f"生成GIF预览，共{len(gif_frames)}帧")
                # 保存GIF，设置较低的帧率使文件更小
                gif_frames[0].save(
                    gif_output, 
                    save_all=True, 
                    append_images=gif_frames[1:], 
                    optimize=True, 
                    duration=500,  # 每帧显示500毫秒
                    loop=0
                )
                gif_created = os.path.exists(gif_output)
                print(f"GIF预览{'成功' if gif_created else '失败'}: {gif_output}")
            except Exception as e:
                print(f"生成GIF预览时出错: {e}")
        
        # 转换为更兼容的格式 - 尝试使用FFmpeg (如果安装了的话)
        try:
            import subprocess
            web_compatible_output = os.path.splitext(output_filename)[0] + '_web.mp4'
            ffmpeg_cmd = [
                'ffmpeg', '-y',
                '-i', output_filename,
                '-vcodec', 'libx264',
                '-pix_fmt', 'yuv420p',
                '-preset', 'fast',
                '-crf', '22',
                web_compatible_output
            ]
            print(f"尝试运行FFmpeg优化视频: {' '.join(ffmpeg_cmd)}")
            
            result = subprocess.run(ffmpeg_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            if result.returncode == 0 and os.path.exists(web_compatible_output):
                print(f"FFmpeg转换成功，使用优化后的视频: {web_compatible_output}")
                output_filename = web_compatible_output
            else:
                print("FFmpeg转换失败，使用原始视频")
        except Exception as e:
            print(f"FFmpeg转换尝试失败: {e}")
        
        # 统计检测结果
        if not all_detections:
            print("注意：未检测到任何物体")
            result_data = {
                'success': True,
                'results': [],
                'video_path': get_web_path(output_filename),
                'message': '未检测到任何水果'
            }
            if gif_created:
                result_data['gif_preview'] = get_web_path(gif_output)
            return result_data
        
        detection_stats = {}
        for det in all_detections:
            cls_name = det['class']
            if cls_name not in detection_stats:
                detection_stats[cls_name] = {
                    'count': 0,
                    'total_confidence': 0
                }
            detection_stats[cls_name]['count'] += 1
            detection_stats[cls_name]['total_confidence'] += det['confidence']
        
        # 计算平均置信度并获取水果信息
        results = []
        for cls_name, stats in detection_stats.items():
            avg_confidence = stats['total_confidence'] / stats['count']
            fruit_info = get_fruit_info(cls_name)
            
            results.append({
                'class': cls_name,
                'chinese_name': fruit_info['name'],
                'confidence': avg_confidence,
                'count': stats['count'],
                'description': fruit_info['description'],
                'nutrition': fruit_info['nutrition'],
                'benefits': fruit_info['benefits'],
                'precautions': fruit_info['precautions']
            })
        
        # 处理输出路径
        web_path = get_web_path(output_filename)
        
        print(f"视频处理完成，输出路径: {output_filename}, Web路径: {web_path}")
        
        result_data = {
            'success': True,
            'results': results,
            'video_path': web_path
        }
        
        # 添加GIF预览路径（如果创建成功）
        if gif_created:
            result_data['gif_preview'] = get_web_path(gif_output)
            print(f"添加GIF预览路径: {get_web_path(gif_output)}")
            
        return result_data
        
    except Exception as e:
        print(f"处理视频时出错: {e}")
        import traceback
        traceback.print_exc()
        return {
            'success': False,
            'error': str(e)
        }

@app.route('/detect_video', methods=['POST'])
def detect_video():
    if 'file' not in request.files:
        return jsonify({'success': False, 'error': '没有上传文件'}), 400
    
    file = request.files['file']
    if file.filename == '':
        return jsonify({'success': False, 'error': '没有选择文件'}), 400
    
    if file and file.filename.lower().endswith(('.mp4', '.avi', '.mov', '.wmv', '.mkv')):
        try:
            # 生成唯一文件名
            filename = f"{datetime.now().strftime('%Y%m%d_%H%M%S')}_{file.filename}"
            filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
            file.save(filepath)
            
            print(f"视频文件已保存到: {filepath}")
            
            # 检查文件大小，过大会导致处理时间过长
            file_size_mb = os.path.getsize(filepath) / (1024 * 1024)
            print(f"视频文件大小: {file_size_mb:.2f} MB")
            
            if file_size_mb > 100:  # 限制100MB
                return jsonify({
                    'success': False, 
                    'error': f'视频文件过大 ({file_size_mb:.2f} MB)，请上传小于100MB的视频'
                })
            
            # 检查文件是否可读
            try:
                check_cap = cv2.VideoCapture(filepath)
                if not check_cap.isOpened():
                    return jsonify({
                        'success': False, 
                        'error': f'无法打开视频文件，可能是不支持的格式或已损坏'
                    })
                
                # 检查视频时长
                fps = check_cap.get(cv2.CAP_PROP_FPS)
                frame_count = check_cap.get(cv2.CAP_PROP_FRAME_COUNT)
                
                if fps <= 0:
                    check_cap.release()
                    return jsonify({
                        'success': False, 
                        'error': f'无效的视频FPS: {fps}'
                    })
                
                duration = frame_count / fps if fps > 0 else 0
                print(f"视频时长: {duration:.2f} 秒")
                
                if duration > 300:  # 限制5分钟
                    check_cap.release()
                    return jsonify({
                        'success': False, 
                        'error': f'视频过长 ({duration:.2f} 秒)，请上传小于5分钟的视频'
                    })
                
                check_cap.release()
            except Exception as e:
                print(f"检查视频时出错: {e}")
                return jsonify({
                    'success': False, 
                    'error': f'检查视频时出错: {str(e)}'
                })
            
            # 加载模型
            try:
                model = load_model()
                if model is None:
                    return jsonify({'success': False, 'error': '模型加载失败'})
            except Exception as e:
                print(f"加载模型时出错: {e}")
                return jsonify({'success': False, 'error': f'加载模型时出错: {str(e)}'})
            
            # 处理视频
            print(f"开始处理视频: {filepath}")
            result = process_video(filepath, model)
            print(f"视频处理结果: {result}")
            
            if not result.get('success', False):
                return jsonify(result)
            
            # 验证输出视频文件是否存在
            video_path = result.get('video_path', '')
            local_path = video_path
            if local_path.startswith('/'):
                local_path = local_path[1:]  # 移除开头的斜杠
            
            if not os.path.exists(local_path) and 'static' in local_path:
                local_path = local_path[local_path.find('static'):]
                
            if not os.path.exists(local_path):
                print(f"警告：输出视频文件不存在: {local_path}")
                result['warning'] = f'输出视频可能无法正常访问: {video_path}'
            
            return jsonify(result)
        except Exception as e:
            print(f"处理视频请求时出错: {e}")
            import traceback
            traceback.print_exc()
            return jsonify({'success': False, 'error': str(e)})
    
    return jsonify({'success': False, 'error': '不支持的文件类型，请上传MP4、AVI或MOV格式的视频'})

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000) 